'use client'
import { Button, Card, Form, Select, Upload, Slider, Space, Progress, Message, Switch } from '@arco-design/web-react'
import { RequestOptions } from '@arco-design/web-react/es/Upload'
import React, { useEffect, useMemo, useRef, useState } from 'react'
import DragDropSubtitle from '../../../components/DragDropSubtitle'
import { NodeList, stringifySync, Node, NodeCue } from 'subtitle'
import { chunkArray, processInParallel } from '../../../utils/parallel'
import { SubtitleResult } from '../../../components/SubtitleResult'

import { clone, cloneDeep, compact, isNil, round } from 'lodash-es'
import { SubTitle } from '../../../types'
import { useTRPC, useTRPCClient } from '../../server/trpc/client'
import { useQuery } from '@tanstack/react-query'
import { useStorageForm } from '../../../hooks/useStorageForm'
import { useLocalStorageState } from 'ahooks'
import { MODELS_STORAGE_KEY } from '../../../constants/models'
import { Model } from '../model-manager/page'

const FormItem = Form.Item
const Option = Select.Option

type FormData = {
  subtiles: NodeList
  chunkSize: number
  maxConcurrent: number
  modelId?: string
  fileName?: string
  enableSummary?: boolean
}
export type SubtitleNode = Node & {
  data: Node['data'] & { transText?: string },

}


const Page = () => {
  const { form, saveForm } = useStorageForm<FormData>({ storageKey: "subtitle_translate", ignoreFields: ["subtiles", "fileName"] })
  const [loading, setLoading] = useState(false);
  const [translatedValue, setTranslatedValue] = useState<SubtitleNode[]>([])
  const [fileName, setFileName] = useState<string>()
  const trpc = useTRPCClient();
  const [models] = useLocalStorageState<Model[]>(MODELS_STORAGE_KEY, { defaultValue: [] });

  const summaryMap = useRef<Record<string, string>>({}); 
  const [pauseFn, setPauseFn] = useState<() => void>()

  

  const process = useMemo(() => {
    const completes = translatedValue.filter(item => item.data.transText !== undefined);
    if (translatedValue.length === 0) return 0;
    
    return completes.length / translatedValue.length
  }, [translatedValue])

  const handleSubmit = async () => {
    form.validate().then(async (res) => {
      setLoading(true);
      saveForm();
      const { subtiles = [], chunkSize = 50, maxConcurrent = 5, enableSummary = true } = form.getFieldsValue();
      const cueItems = subtiles.filter(item => item.type === 'cue');
      const strs = cueItems.map((item, index) => ({
        index,
        text: item.data.text,
      }));
      
      let summary = '';
      if (enableSummary) {
        summary = await trpc.subtile.summarizeSubtitle.mutate({
          content: strs.map(item => item.text).join('\n'),
          modelConfig: getModelInfo(),
        });
        summaryMap.current[fileName!] = summary

        console.log('总结内容:', summary);
      }
      
      setTranslatedValue(cloneDeep(cueItems));
      handleTranslate(strs, chunkSize, maxConcurrent, summary);
    })

  }
  console.log(summaryMap, fileName);
  
  const handleContinue = async () => {
    setLoading(true);

    const { chunkSize = 50, maxConcurrent = 5 } = form.getFieldsValue();

    const strs = translatedValue.map((item, index) => {
      if (item.data.transText === undefined) {
        return {
          index,
          text: (item as NodeCue).data.text,
        }
      }
      return undefined
    });
    console.log(strs);
    

    handleTranslate(compact(strs), chunkSize, maxConcurrent, summaryMap.current[fileName!]);

  }


  const handleTranslate = async (subtiles: SubTitle[], chunkSize = 50, maxConcurrent = 5, summary = '') => {

    setLoading(true);
    const chunks = chunkArray(subtiles, chunkSize);

    try {

      const { pause , promise } = processInParallel(
        chunks,
        chunk => handleRemoteTrans(chunk, summary),
        maxConcurrent,
      );

      setPauseFn(() => pause)

      await promise;

      // setTanslatedValue(translatedSubtitles);
    } catch (error) {
      
      Message.error(`翻译出错:${error}`);
      throw error;
    } finally {
      setLoading(false);
      setPauseFn(undefined);
    }
  }

  const handlePause = () => {
    setLoading(false);
    pauseFn?.()
    setPauseFn(undefined);

    
  }

  const getModelInfo = () => {
    const modelId = form.getFieldValue('modelId')
    const selectedModel = models.find(m => m.id === modelId);
    return selectedModel!
  }

  const handleRemoteTrans = (chunk: SubTitle[], summary: string) => {

    return trpc.subtile.translateSubtitle.mutate({
      subtitles: chunk,
      modelConfig: getModelInfo(),
      summary
    }).then(async iterable => {
      return new Promise<SubTitle[]>(async (res, rej) => {
        const acc: SubTitle[] = []
        console.log(11);
        
        try {
          for await (const value of iterable) {
            const { type, data } = value
            switch(type) {
            case 'data':
              if (data) {
                acc.push(data);
                
                setTranslatedValue(newItems => {
                  const results = cloneDeep(newItems);
      
                  const { index, text } = data;
                  const val = results[index]
                  if (val) {
                    val.data.transText = text;
                  }
                  console.log('翻译返回结果处理完成', data);
              
                  return results;
                })
              }
              break;
            case 'end':
              res(acc);
              break;
            default:
              break;
            }

          }
        
        } catch (e) {
          rej(e);
        } finally {
          res(acc);
          
        }
      })
    })
  }

  const generateSrtContent = () => {
    if (translatedValue.length === 0) return [];
    
    // 过滤掉transText为"NULL"的条目
    const filteredSubtitles = translatedValue.filter(
      item => item.data.transText !== null
    );
    
    // 将transText应用到text字段
    const subtitlesForExport = filteredSubtitles.map(item => {
      if (item.type === 'cue' && item.data.transText) {
        item.data.text = item.data.transText;
      }
      return item;
    });
    return subtitlesForExport;
  }

  const downloadSrt = () => {
    const subtitlesForExport = generateSrtContent();
    const srtContent = stringifySync(subtitlesForExport, { format: 'SRT' });
    const currentFileName = (typeof form.getFieldValue('fileName') === 'string' 
      ? form.getFieldValue('fileName') 
      : fileName) as string | undefined;
    downloadFile(srtContent, currentFileName);
  }

  const downloadFile= (content: string, originalName?: string) => {
    const blob = new Blob([content], { type: 'text/plain' });
    const url = URL.createObjectURL(blob);
    const a = document.createElement('a');
    a.href = url;
    a.download = originalName 
      ? originalName.replace('.srt', '.trans.srt') 
      : 'translated_subtitles.srt';
    document.body.appendChild(a);
    a.click();
    document.body.removeChild(a);
    URL.revokeObjectURL(url);
  }

  const copySubtitles = () => {
    const subtitlesForExport = generateSrtContent();
    
    const srtContent = stringifySync(subtitlesForExport, { format: 'SRT' });
    navigator.clipboard.writeText(srtContent)
      .then(() => Message.success("复制成功"))
      .catch(err =>  Message.error(`复制失败: ${err}`));
  }

  return (
    <Card title="字幕翻译" className={"max-w-200"}>
      <Form layout="vertical" form={form} onSubmit={handleSubmit}>
        {models.length > 0 && (
          <Form.Item label="选择模型" field="modelId" rules={[{ required: true }]}>
            <Select>
              {models.map(model => (
                <Select.Option key={model.id} value={model.id}>
                  {model.name} ({model.model})
                </Select.Option>
              ))}
            </Select>
          </Form.Item>
        )}
        <FormItem rules={[{ required: true }]} required label='上传字幕文件' field="subtiles">
          <DragDropSubtitle onChange={(nodes, name) => {
            form.setFieldValue('subtiles', nodes);
            form.setFieldValue('fileName', name);
            setFileName(name);
          }} />
        </FormItem>
        <FormItem label="每批翻译字幕行数（越多对上下文识别越好）" field="chunkSize" initialValue={50}>
          <Slider min={1} max={500} step={1} showInput />
        </FormItem>
        <FormItem label="最大并行处理批数（越多越快）" field="maxConcurrent" initialValue={5}>
          <Slider min={1} max={20} step={1} showInput />
        </FormItem>
        <FormItem label="翻译前总结内容（有助于上下文理解）"triggerPropName="checked" field="enableSummary" initialValue={true}>
          <Switch />
        </FormItem>
        <FormItem>
          <Space direction='vertical' className={"w-full"}>
            {!pauseFn && <Button shape="round" long loading={loading} size="large" type='primary' onClick={handleSubmit}>翻译</Button>}
            {pauseFn && <Button shape="round" long size="large" type='primary' status="danger" onClick={handlePause}>暂停</Button>}
            {!pauseFn && ![0, 1].includes(process) && <Button shape="round" long loading={loading} size="large" type='primary' status="success" onClick={handleContinue}>翻译剩余部分</Button>}
          </Space>
        </FormItem>
      </Form>
      <Progress percent={round(process * 100, 2)}/>
      <SubtitleResult 
        cardProps={{
          extra: (
            (process != 0 && <Space>
              <Button 
                size='mini' 
                type="primary"
                onClick={downloadSrt}
              >
                下载翻译字幕
              </Button>
              <Button 
                size='mini' 
                type="primary"
                onClick={copySubtitles}
              >
                复制字幕
              </Button>
            </Space>)
          )
        }} 
        items={translatedValue} 
      />
    </Card>
  )
}

export default Page